import numpy as np
import matplotlib.pyplot as plt

class TrajectoryGenerator:
    """
    轨迹生成器
    """
    def __init__(self, dt=0.1):
        self.dt = dt
        
    def circle_trajectory(self, t, radius=5, center=(0, 0), omega=0.2):
        """
        生成圆形轨迹
        t: 时间
        radius: 半径
        center: 圆心坐标
        omega: 角速度
        """
        x_c, y_c = center
        
        # 位置
        x = x_c + radius * np.cos(omega * t)
        y = y_c + radius * np.sin(omega * t)
        
        # 速度
        vx = -radius * omega * np.sin(omega * t)
        vy = radius * omega * np.cos(omega * t)
        
        return np.array([x, y, vx, vy])
    
    def figure_eight_trajectory(self, t, scale=5, omega=0.2):
        """
        生成8字形轨迹
        t: 时间
        scale: 尺度因子
        omega: 角速度
        """
        # 位置
        x = scale * np.sin(omega * t)
        y = scale * np.sin(omega * t) * np.cos(omega * t)
        
        # 速度
        vx = scale * omega * np.cos(omega * t)
        vy = scale * omega * (np.cos(omega * t) * np.cos(omega * t) - np.sin(omega * t) * np.sin(omega * t))
        
        return np.array([x, y, vx, vy])
    
    def square_trajectory(self, t, side_length=5, period=20):
        """
        生成方形轨迹
        t: 时间
        side_length: 边长
        period: 周期
        """
        # 一个周期内的时间
        t_mod = t % period
        half_side = side_length / 2
        
        # 分段生成轨迹
        if t_mod < period/4:  # 上边
            x = -half_side + side_length * (4 * t_mod / period)
            y = half_side
            vx = side_length * (4 / period)
            vy = 0
        elif t_mod < period/2:  # 右边
            x = half_side
            y = half_side - side_length * (4 * (t_mod - period/4) / period)
            vx = 0
            vy = -side_length * (4 / period)
        elif t_mod < 3*period/4:  # 下边
            x = half_side - side_length * (4 * (t_mod - period/2) / period)
            y = -half_side
            vx = -side_length * (4 / period)
            vy = 0
        else:  # 左边
            x = -half_side
            y = -half_side + side_length * (4 * (t_mod - 3*period/4) / period)
            vx = 0
            vy = side_length * (4 / period)
        
        return np.array([x, y, vx, vy])
    
    def generate_trajectory(self, trajectory_type, times):
        """
        根据轨迹类型生成完整轨迹
        trajectory_type: 轨迹类型 ('circle', 'figure_eight', 'square')
        times: 时间数组
        """
        trajectory = []
        
        # 设置合适的轨迹参数
        if trajectory_type == 'circle':
            # 使用小半径的圆形轨迹，确保它清晰可见
            radius = 5
            center = (0, 0)  # 使用原点作为圆心
            omega = 0.5  # 增加角速度使轨迹更明显
            
            print(f"生成圆形轨迹：半径={radius}, 圆心={center}, 角速度={omega}")
            
            for t in times:
                state = self.circle_trajectory(t, radius=radius, center=center, omega=omega)
                trajectory.append(state)
                
        elif trajectory_type == 'figure_eight':
            scale = 5
            omega = 0.3
            
            for t in times:
                state = self.figure_eight_trajectory(t, scale=scale, omega=omega)
                trajectory.append(state)
                
        elif trajectory_type == 'square':
            side_length = 5
            period = 20
            
            for t in times:
                state = self.square_trajectory(t, side_length=side_length, period=period)
                trajectory.append(state)
                
        else:
            raise ValueError(f"未知轨迹类型: {trajectory_type}")
        
        # 转换为numpy数组
        trajectory_array = np.array(trajectory)
        
        # 检查轨迹数据的有效性
        if np.any(np.isnan(trajectory_array)) or np.any(np.isinf(trajectory_array)):
            raise ValueError("轨迹生成出现NaN或Inf值，请检查参数设置")
            
        # 打印轨迹信息
        print(f"生成{trajectory_type}轨迹: {len(trajectory_array)}个点")
        print(f"轨迹范围: X [{np.min(trajectory_array[:,0]):.2f}, {np.max(trajectory_array[:,0]):.2f}], "
              f"Y [{np.min(trajectory_array[:,1]):.2f}, {np.max(trajectory_array[:,1]):.2f}]")
        
        # 可视化生成的轨迹（仅用于调试）
        self._visualize_trajectory(trajectory_array, trajectory_type)
        
        return trajectory_array
        
    def _visualize_trajectory(self, trajectory, trajectory_type):
        """可视化生成的轨迹（仅用于调试）"""
        plt.figure(figsize=(8, 8))
        plt.plot(trajectory[:, 0], trajectory[:, 1], 'b-')
        plt.scatter(trajectory[0, 0], trajectory[0, 1], c='g', s=100, label='起点')
        
        # 绘制每100个点一个标记，显示方向
        step = 100
        for i in range(0, len(trajectory), step):
            if i+step < len(trajectory):
                plt.arrow(trajectory[i, 0], trajectory[i, 1], 
                         (trajectory[i+step, 0] - trajectory[i, 0])/10, 
                         (trajectory[i+step, 1] - trajectory[i, 1])/10,
                         head_width=0.2, head_length=0.3, fc='r', ec='r')
        
        plt.title(f'{trajectory_type}参考轨迹')
        plt.xlabel('X位置')
        plt.ylabel('Y位置')
        plt.grid(True)
        plt.axis('equal')
        plt.legend()
        
        # 保存到文件
        plt.savefig(f'results/{trajectory_type}_reference.png', dpi=200, bbox_inches='tight')
        
        # 不在这里显示，只保存文件
        plt.close()
        
        print(f"参考轨迹已保存到: results/{trajectory_type}_reference.png")
        
    def test_trajectories(self):
        """测试生成所有轨迹类型并可视化"""
        import os
        
        # 创建results目录（如果不存在）
        if not os.path.exists('results'):
            os.makedirs('results')
            
        # 生成时间点
        times = np.arange(0, 30, 0.1)
        
        # 生成并可视化所有轨迹类型
        for traj_type in ['circle', 'figure_eight', 'square']:
            trajectory = self.generate_trajectory(traj_type, times)
            
            # 在这里直接显示图像
            plt.figure(figsize=(10, 8))
            plt.plot(trajectory[:, 0], trajectory[:, 1], 'b-', label='轨迹')
            plt.scatter(trajectory[0, 0], trajectory[0, 1], c='g', s=100, label='起点')
            plt.title(f'{traj_type}参考轨迹')
            plt.xlabel('X位置')
            plt.ylabel('Y位置')
            plt.grid(True)
            plt.axis('equal')
            plt.legend()
            plt.show()
            
        return True 